# -*- coding: utf-8 -*-
"""
管理员管理模块 - 认证与安全组件

本文件实现了测试DEMO中管理员认证和安全相关的核心组件，包括：

1. JWTAuthentication - JWT认证类
   - 基于JWT令牌的用户认证
   - 验证管理员身份和权限
   - 处理令牌过期和无效情况

2. JWTTokenGenerator - JWT令牌生成器
   - generate_token: 生成访问令牌
   - generate_refresh_token: 生成刷新令牌
   - verify_refresh_token: 验证刷新令牌

3. TwoFactorAuth - 双因素认证
   - generate_secret: 生成TOTP密钥
   - generate_qr_code: 生成二维码
   - verify_totp: 验证TOTP令牌
   - generate_backup_codes: 生成备用验证码

4. SessionManager - 会话管理器
   - create_session: 创建用户会话
   - validate_session: 验证会话有效性
   - invalidate_session: 使会话失效
   - cleanup_expired_sessions: 清理过期会话

5. PasswordSecurity - 密码安全
   - hash_password: 密码哈希加密
   - verify_password: 密码验证
   - check_password_strength: 密码强度检查
   - generate_secure_password: 生成安全密码

6. SecurityLogger - 安全日志记录器
   - log_login_attempt: 记录登录尝试
   - log_password_change: 记录密码更改
   - log_permission_denied: 记录权限拒绝
   - log_suspicious_activity: 记录可疑活动

7. IPWhitelist - IP白名单管理
   - is_ip_allowed: 检查IP是否在白名单
   - add_ip_to_whitelist: 添加IP到白名单
   - remove_ip_from_whitelist: 从白名单移除IP

这些组件共同构成了完整的安全认证体系，确保系统的安全性和可靠性。

Author: zhaozeliang
Date: 2025-09-06
"""

import jwt  # type: ignore
import datetime
from django.conf import settings  # type: ignore
from django.contrib.auth import authenticate  # type: ignore
from rest_framework import authentication, exceptions  # type: ignore
from rest_framework.authtoken.models import Token  # type: ignore
from .models import AdminUser
import hashlib
import secrets
import pyotp  # type: ignore
import qrcode
from io import BytesIO
import base64


class JWTAuthentication(authentication.BaseAuthentication):
    """
    JWT认证类
    """
    
    def authenticate(self, request):
        """
        认证用户
        """
        auth_header = request.META.get('HTTP_AUTHORIZATION')
        
        if not auth_header or not auth_header.startswith('Bearer '):
            return None
        
        token = auth_header.split(' ')[1]
        
        try:
            payload = jwt.decode(
                token,
                settings.SECRET_KEY,
                algorithms=['HS256']
            )
        except jwt.ExpiredSignatureError:
            raise exceptions.AuthenticationFailed('Token已过期')
        except jwt.InvalidTokenError:
            raise exceptions.AuthenticationFailed('无效的Token')
        
        user_id = payload.get('user_id')
        if not user_id:
            raise exceptions.AuthenticationFailed('Token中缺少用户ID')
        
        try:
            from django.contrib.auth.models import User  # type: ignore  # type: ignore
            user = User.objects.get(id=user_id)
        except User.DoesNotExist:
            raise exceptions.AuthenticationFailed('用户不存在')
        
        # 检查用户是否为管理员
        try:
            admin_user = AdminUser.objects.get(user=user)
            if not admin_user.is_active:
                raise exceptions.AuthenticationFailed('管理员账户已被禁用')
        except AdminUser.DoesNotExist:
            raise exceptions.AuthenticationFailed('用户不是管理员')
        
        return (user, token)


class JWTTokenGenerator:
    """
    JWT Token生成器
    """
    
    @staticmethod
    def generate_token(user, expires_in_hours=24):
        """
        生成JWT Token
        """
        payload = {
            'user_id': user.id,
            'username': user.username,
            'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=expires_in_hours),
            'iat': datetime.datetime.utcnow(),
        }
        
        token = jwt.encode(
            payload,
            settings.SECRET_KEY,
            algorithm='HS256'
        )
        
        return token
    
    @staticmethod
    def generate_refresh_token(user):
        """
        生成刷新Token
        """
        payload = {
            'user_id': user.id,
            'type': 'refresh',
            'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
            'iat': datetime.datetime.utcnow(),
        }
        
        token = jwt.encode(
            payload,
            settings.SECRET_KEY,
            algorithm='HS256'
        )
        
        return token
    
    @staticmethod
    def verify_refresh_token(token):
        """
        验证刷新Token
        """
        try:
            payload = jwt.decode(
                token,
                settings.SECRET_KEY,
                algorithms=['HS256']
            )
            
            if payload.get('type') != 'refresh':
                return None
            
            user_id = payload.get('user_id')
            if not user_id:
                return None
            
            from django.contrib.auth.models import User
            user = User.objects.get(id=user_id)
            return user
        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError, User.DoesNotExist):
            return None


class TwoFactorAuth:
    """
    双因素认证
    """
    
    @staticmethod
    def generate_secret():
        """
        生成TOTP密钥
        """
        return pyotp.random_base32()
    
    @staticmethod
    def generate_qr_code(user, secret):
        """
        生成二维码
        """
        totp_uri = pyotp.totp.TOTP(secret).provisioning_uri(
            name=user.username,
            issuer_name="测试DEMO"
        )
        
        qr = qrcode.QRCode(version=1, box_size=10, border=5)
        qr.add_data(totp_uri)
        qr.make(fit=True)
        
        img = qr.make_image(fill_color="black", back_color="white")
        
        # 转换为base64
        buffer = BytesIO()
        img.save(buffer, format='PNG')
        img_str = base64.b64encode(buffer.getvalue()).decode()
        
        return f"data:image/png;base64,{img_str}"
    
    @staticmethod
    def verify_totp(secret, token):
        """
        验证TOTP令牌
        """
        totp = pyotp.TOTP(secret)
        return totp.verify(token, valid_window=1)
    
    @staticmethod
    def generate_backup_codes(count=10):
        """
        生成备用代码
        """
        codes = []
        for _ in range(count):
            code = secrets.token_hex(4).upper()
            codes.append(f"{code[:4]}-{code[4:]}")
        return codes


class SessionManager:
    """
    会话管理器
    """
    
    @staticmethod
    def create_session(user, ip_address, user_agent):
        """
        创建会话
        """
        from .models import UserSession
        
        # 生成会话ID
        session_id = secrets.token_urlsafe(32)
        
        # 创建会话记录
        session = UserSession.objects.create(
            user=user,
            session_id=session_id,
            ip_address=ip_address,
            user_agent=user_agent,
            expires_at=datetime.datetime.now() + datetime.timedelta(hours=24)
        )
        
        return session
    
    @staticmethod
    def validate_session(session_id, user):
        """
        验证会话
        """
        from .models import UserSession
        
        try:
            session = UserSession.objects.get(
                session_id=session_id,
                user=user,
                is_active=True,
                expires_at__gt=datetime.datetime.now()
            )
            
            # 更新最后活动时间
            session.last_activity = datetime.datetime.now()
            session.save()
            
            return session
        except UserSession.DoesNotExist:
            return None
    
    @staticmethod
    def invalidate_session(session_id):
        """
        使会话失效
        """
        from .models import UserSession
        
        UserSession.objects.filter(session_id=session_id).update(is_active=False)
    
    @staticmethod
    def cleanup_expired_sessions():
        """
        清理过期会话
        """
        from .models import UserSession
        
        UserSession.objects.filter(
            expires_at__lt=datetime.datetime.now()
        ).delete()


class PasswordSecurity:
    """
    密码安全工具
    """
    
    @staticmethod
    def hash_password(password, salt=None):
        """
        哈希密码
        """
        if salt is None:
            salt = secrets.token_hex(16)
        
        # 使用PBKDF2进行密码哈希
        password_hash = hashlib.pbkdf2_hmac(
            'sha256',
            password.encode('utf-8'),
            salt.encode('utf-8'),
            100000  # 迭代次数
        )
        
        return {
            'hash': password_hash.hex(),
            'salt': salt
        }
    
    @staticmethod
    def verify_password(password, stored_hash, salt):
        """
        验证密码
        """
        password_hash = hashlib.pbkdf2_hmac(
            'sha256',
            password.encode('utf-8'),
            salt.encode('utf-8'),
            100000
        )
        
        return password_hash.hex() == stored_hash
    
    @staticmethod
    def check_password_strength(password):
        """
        检查密码强度
        """
        import re
        
        errors = []
        
        # 长度检查
        if len(password) < 8:
            errors.append("密码长度至少8位")
        
        # 包含大写字母
        if not re.search(r'[A-Z]', password):
            errors.append("密码必须包含大写字母")
        
        # 包含小写字母
        if not re.search(r'[a-z]', password):
            errors.append("密码必须包含小写字母")
        
        # 包含数字
        if not re.search(r'\d', password):
            errors.append("密码必须包含数字")
        
        # 包含特殊字符
        if not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
            errors.append("密码必须包含特殊字符")
        
        return {
            'is_strong': len(errors) == 0,
            'errors': errors
        }
    
    @staticmethod
    def generate_secure_password(length=12):
        """
        生成安全密码
        """
        import string
        
        characters = string.ascii_letters + string.digits + "!@#$%^&*"
        password = ''.join(secrets.choice(characters) for _ in range(length))
        
        # 确保包含各种字符类型
        if not any(c.isupper() for c in password):
            password = password[:-1] + secrets.choice(string.ascii_uppercase)
        
        if not any(c.islower() for c in password):
            password = password[:-1] + secrets.choice(string.ascii_lowercase)
        
        if not any(c.isdigit() for c in password):
            password = password[:-1] + secrets.choice(string.digits)
        
        if not any(c in "!@#$%^&*" for c in password):
            password = password[:-1] + secrets.choice("!@#$%^&*")
        
        return password


class SecurityLogger:
    """
    安全日志记录器
    """
    
    @staticmethod
    def log_login_attempt(user, ip_address, success, reason=None):
        """
        记录登录尝试
        """
        from .models import SecurityLog
        
        SecurityLog.objects.create(
            user=user,
            event_type='login_attempt',
            ip_address=ip_address,
            success=success,
            details={
                'reason': reason
            }
        )
    
    @staticmethod
    def log_password_change(user, ip_address):
        """
        记录密码更改
        """
        from .models import SecurityLog
        
        SecurityLog.objects.create(
            user=user,
            event_type='password_change',
            ip_address=ip_address,
            success=True
        )
    
    @staticmethod
    def log_permission_denied(user, resource, ip_address):
        """
        记录权限拒绝
        """
        from .models import SecurityLog
        
        SecurityLog.objects.create(
            user=user,
            event_type='permission_denied',
            ip_address=ip_address,
            success=False,
            details={
                'resource': resource
            }
        )
    
    @staticmethod
    def log_suspicious_activity(user, activity, ip_address, details=None):
        """
        记录可疑活动
        """
        from .models import SecurityLog
        
        SecurityLog.objects.create(
            user=user,
            event_type='suspicious_activity',
            ip_address=ip_address,
            success=False,
            details={
                'activity': activity,
                'details': details or {}
            }
        )


class IPWhitelist:
    """
    IP白名单管理
    """
    
    @staticmethod
    def is_ip_allowed(ip_address):
        """
        检查IP是否在白名单中
        """
        from .models import IPWhitelist as IPWhitelistModel
        
        return IPWhitelistModel.objects.filter(
            ip_address=ip_address,
            is_active=True
        ).exists()
    
    @staticmethod
    def add_ip_to_whitelist(ip_address, description=""):
        """
        添加IP到白名单
        """
        from .models import IPWhitelist as IPWhitelistModel
        
        whitelist_entry, created = IPWhitelistModel.objects.get_or_create(
            ip_address=ip_address,
            defaults={
                'description': description,
                'is_active': True
            }
        )
        
        return whitelist_entry
    
    @staticmethod
    def remove_ip_from_whitelist(ip_address):
        """
        从白名单中移除IP
        """
        from .models import IPWhitelist as IPWhitelistModel
        
        IPWhitelistModel.objects.filter(ip_address=ip_address).delete()